{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "from torch.optim import lr_scheduler\n",
    "import numpy as np\n",
    "import time\n",
    "import copy\n",
    "import os\n",
    "class h_sigmoid(nn.Module):\n",
    "    def __init__(self, inplace=True):\n",
    "        super(h_sigmoid, self).__init__()\n",
    "        self.relu = nn.ReLU6(inplace=inplace)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.relu(x + 3) / 6\n",
    "\n",
    "\n",
    "class h_swish(nn.Module):\n",
    "    def __init__(self, inplace=True):\n",
    "        super(h_swish, self).__init__()\n",
    "        self.sigmoid = h_sigmoid(inplace=inplace)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x * self.sigmoid(x)\n",
    "\n",
    "\n",
    "class CoordAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, reduction=32):\n",
    "        super(CoordAttention, self).__init__()\n",
    "        self.pool_w, self.pool_h = nn.AdaptiveAvgPool2d((1, None)), nn.AdaptiveAvgPool2d((None, 1))\n",
    "        temp_c = max(8, in_channels // reduction)\n",
    "        self.conv1 = nn.Conv2d(in_channels, temp_c, kernel_size=1, stride=1, padding=0)\n",
    "\n",
    "        self.bn1 = nn.BatchNorm2d(temp_c)\n",
    "        self.act1 = h_swish()\n",
    "\n",
    "        self.conv2 = nn.Conv2d(temp_c, out_channels, kernel_size=1, stride=1, padding=0)\n",
    "        self.conv3 = nn.Conv2d(temp_c, out_channels, kernel_size=1, stride=1, padding=0)\n",
    "\n",
    "    def forward(self, x):\n",
    "        short = x\n",
    "        n, c, H, W = x.shape\n",
    "        x_h, x_w = self.pool_h(x), self.pool_w(x).permute(0, 1, 3, 2)\n",
    "        x_cat = torch.cat([x_h, x_w], dim=2)\n",
    "        out = self.act1(self.bn1(self.conv1(x_cat)))\n",
    "        x_h, x_w = torch.split(out, [H, W], dim=2)\n",
    "        x_w = x_w.permute(0, 1, 3, 2)\n",
    "        out_h = torch.sigmoid(self.conv2(x_h))\n",
    "        out_w = torch.sigmoid(self.conv3(x_w))\n",
    "        return short * out_w * out_h\n",
    "\n",
    "# 定义基本模块\n",
    "class BasicBlock(nn.Module):\n",
    "    expansion = 1\n",
    "\n",
    "    def __init__(self, in_planes, planes, stride=1):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != self.expansion * planes:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(self.expansion * planes)\n",
    "            )\n",
    "        \n",
    "        # 添加 CoordAttention\n",
    "        self.coord_attention = CoordAttention(planes, planes)  # 调整输入输出通道数\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.bn2(self.conv2(out))\n",
    "        out += self.shortcut(x)\n",
    "        out = F.relu(out)\n",
    "        \n",
    "        # 使用 CoordAttention\n",
    "        out = self.coord_attention(out)\n",
    "        return out\n",
    "\n",
    "# 定义 ResNet\n",
    "class ResNet(nn.Module):\n",
    "    def __init__(self, block, num_blocks, num_classes=100):\n",
    "        super(ResNet, self).__init__()\n",
    "        self.in_planes = 64\n",
    "\n",
    "        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(64)\n",
    "        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)\n",
    "        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)\n",
    "        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)\n",
    "        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)\n",
    "        self.linear = nn.Linear(512 * block.expansion, num_classes)\n",
    "\n",
    "    def _make_layer(self, block, planes, num_blocks, stride):\n",
    "        strides = [stride] + [1] * (num_blocks - 1)\n",
    "        layers = []\n",
    "        for stride in strides:\n",
    "            layers.append(block(self.in_planes, planes, stride))\n",
    "            self.in_planes = planes * block.expansion\n",
    "        \n",
    "        # 在每个 BasicBlock 之后添加 CoordAttention\n",
    "        for i in range(len(layers)):\n",
    "            layers.insert(2*i + 1, CoordAttention(planes, planes))  # 在每个 BasicBlock 之后插入 CoordAttention\n",
    "        \n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.bn1(self.conv1(x)))\n",
    "        out = self.layer1(out)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "        out = self.layer4(out)\n",
    "        out = F.avg_pool2d(out, 4)\n",
    "        out = out.view(out.size(0), -1)\n",
    "        out = self.linear(out)\n",
    "        return out\n",
    "\n",
    "# 定义 ResNet-18\n",
    "def ResNet18():\n",
    "    return ResNet(BasicBlock, [2, 2, 2, 2])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# 数据预处理\n",
    "transform_train = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomCrop(32, padding=4),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5071, 0.4865, 0.4409), (0.2673, 0.2564, 0.2762)),\n",
    "])\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5071, 0.4865, 0.4409), (0.2673, 0.2564, 0.2762)),\n",
    "])\n",
    "\n",
    "# 加载数据集\n",
    "trainset = torchvision.datasets.CIFAR100(root='./data', train=True, download=True, transform=transform_train)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)\n",
    "\n",
    "testset = torchvision.datasets.CIFAR100(root='./data', train=False, download=True, transform=transform_test)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)\n",
    "\n",
    "# 实例化模型\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    "net = ResNet18().to(device)\n",
    "\n",
    "# 定义损失函数和优化器\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "第 0/24 轮训练\n",
      "----------\n",
      "train 损失: 4.5827 Top-1 准确率: 0.0142 Top-5 准确率: 0.0603\n",
      "val 损失: 4.5562 Top-1 准确率: 0.0211 Top-5 准确率: 0.0738\n",
      "\n",
      "第 1/24 轮训练\n",
      "----------\n",
      "train 损失: 4.5414 Top-1 准确率: 0.0214 Top-5 准确率: 0.0780\n",
      "val 损失: 4.5268 Top-1 准确率: 0.0268 Top-5 准确率: 0.0958\n",
      "\n",
      "第 2/24 轮训练\n",
      "----------\n",
      "train 损失: 4.5092 Top-1 准确率: 0.0268 Top-5 准确率: 0.0948\n",
      "val 损失: 4.5150 Top-1 准确率: 0.0321 Top-5 准确率: 0.1072\n",
      "\n",
      "第 3/24 轮训练\n",
      "----------\n",
      "train 损失: 4.4520 Top-1 准确率: 0.0417 Top-5 准确率: 0.1367\n",
      "val 损失: 4.3848 Top-1 准确率: 0.0511 Top-5 准确率: 0.1734\n",
      "\n",
      "第 4/24 轮训练\n",
      "----------\n",
      "train 损失: 4.3183 Top-1 准确率: 0.0561 Top-5 准确率: 0.1875\n",
      "val 损失: 4.2233 Top-1 准确率: 0.0643 Top-5 准确率: 0.2082\n",
      "\n",
      "第 5/24 轮训练\n",
      "----------\n",
      "train 损失: 4.1627 Top-1 准确率: 0.0706 Top-5 准确率: 0.2267\n",
      "val 损失: 4.0872 Top-1 准确率: 0.0720 Top-5 准确率: 0.2367\n",
      "\n",
      "第 6/24 轮训练\n",
      "----------\n",
      "train 损失: 3.9854 Top-1 准确率: 0.0975 Top-5 准确率: 0.2797\n",
      "val 损失: 3.8612 Top-1 准确率: 0.1142 Top-5 准确率: 0.3163\n",
      "\n",
      "第 7/24 轮训练\n",
      "----------\n",
      "train 损失: 3.7631 Top-1 准确率: 0.1338 Top-5 准确率: 0.3522\n",
      "val 损失: 3.6270 Top-1 准确率: 0.1505 Top-5 准确率: 0.3855\n",
      "\n",
      "第 8/24 轮训练\n",
      "----------\n",
      "train 损失: 3.5047 Top-1 准确率: 0.1740 Top-5 准确率: 0.4252\n",
      "val 损失: 3.3256 Top-1 准确率: 0.1989 Top-5 准确率: 0.4745\n",
      "\n",
      "第 9/24 轮训练\n",
      "----------\n",
      "train 损失: 3.1951 Top-1 准确率: 0.2204 Top-5 准确率: 0.5053\n",
      "val 损失: 3.1248 Top-1 准确率: 0.2248 Top-5 准确率: 0.5124\n",
      "\n",
      "第 10/24 轮训练\n",
      "----------\n",
      "train 损失: 2.8774 Top-1 准确率: 0.2767 Top-5 准确率: 0.5832\n",
      "val 损失: 2.7402 Top-1 准确率: 0.2913 Top-5 准确率: 0.6085\n",
      "\n",
      "第 11/24 轮训练\n",
      "----------\n",
      "train 损失: 2.5494 Top-1 准确率: 0.3402 Top-5 准确率: 0.6581\n",
      "val 损失: 2.4538 Top-1 准确率: 0.3592 Top-5 准确率: 0.6781\n",
      "\n",
      "第 12/24 轮训练\n",
      "----------\n",
      "train 损失: 2.2667 Top-1 准确率: 0.3956 Top-5 准确率: 0.7206\n",
      "val 损失: 2.2666 Top-1 准确率: 0.3997 Top-5 准确率: 0.7177\n",
      "\n",
      "第 13/24 轮训练\n",
      "----------\n",
      "train 损失: 2.0330 Top-1 准确率: 0.4492 Top-5 准确率: 0.7690\n",
      "val 损失: 2.1108 Top-1 准确率: 0.4377 Top-5 准确率: 0.7456\n",
      "\n",
      "第 14/24 轮训练\n",
      "----------\n",
      "train 损失: 1.8553 Top-1 准确率: 0.4872 Top-5 准确率: 0.8021\n",
      "val 损失: 2.0741 Top-1 准确率: 0.4500 Top-5 准确率: 0.7525\n",
      "\n",
      "第 15/24 轮训练\n",
      "----------\n",
      "train 损失: 1.6955 Top-1 准确率: 0.5281 Top-5 准确率: 0.8300\n",
      "val 损失: 1.9083 Top-1 准确率: 0.4868 Top-5 准确率: 0.7862\n",
      "\n",
      "第 16/24 轮训练\n",
      "----------\n",
      "train 损失: 1.5701 Top-1 准确率: 0.5563 Top-5 准确率: 0.8496\n",
      "val 损失: 1.7270 Top-1 准确率: 0.5252 Top-5 准确率: 0.8198\n",
      "\n",
      "第 17/24 轮训练\n",
      "----------\n",
      "train 损失: 1.4614 Top-1 准确率: 0.5834 Top-5 准确率: 0.8668\n",
      "val 损失: 1.6769 Top-1 准确率: 0.5357 Top-5 准确率: 0.8324\n",
      "\n",
      "第 18/24 轮训练\n",
      "----------\n",
      "train 损失: 1.3694 Top-1 准确率: 0.6091 Top-5 准确率: 0.8817\n",
      "val 损失: 1.7110 Top-1 准确率: 0.5366 Top-5 准确率: 0.8252\n",
      "\n",
      "第 19/24 轮训练\n",
      "----------\n",
      "train 损失: 1.2736 Top-1 准确率: 0.6332 Top-5 准确率: 0.8965\n",
      "val 损失: 1.5168 Top-1 准确率: 0.5819 Top-5 准确率: 0.8502\n",
      "\n",
      "第 20/24 轮训练\n",
      "----------\n",
      "train 损失: 1.2067 Top-1 准确率: 0.6486 Top-5 准确率: 0.9052\n",
      "val 损失: 1.5112 Top-1 准确率: 0.5772 Top-5 准确率: 0.8606\n",
      "\n",
      "第 21/24 轮训练\n",
      "----------\n",
      "train 损失: 1.1357 Top-1 准确率: 0.6696 Top-5 准确率: 0.9155\n",
      "val 损失: 1.5328 Top-1 准确率: 0.5804 Top-5 准确率: 0.8577\n",
      "\n",
      "第 22/24 轮训练\n",
      "----------\n",
      "train 损失: 1.0747 Top-1 准确率: 0.6841 Top-5 准确率: 0.9238\n",
      "val 损失: 1.4583 Top-1 准确率: 0.5977 Top-5 准确率: 0.8673\n",
      "\n",
      "第 23/24 轮训练\n",
      "----------\n",
      "train 损失: 1.0084 Top-1 准确率: 0.7014 Top-5 准确率: 0.9322\n",
      "val 损失: 1.3987 Top-1 准确率: 0.6094 Top-5 准确率: 0.8754\n",
      "\n",
      "第 24/24 轮训练\n",
      "----------\n",
      "train 损失: 0.9555 Top-1 准确率: 0.7161 Top-5 准确率: 0.9373\n",
      "val 损失: 1.4801 Top-1 准确率: 0.6000 Top-5 准确率: 0.8630\n",
      "\n",
      "训练完成，耗时 44m 44s\n",
      "最佳验证集准确率: 0.609400\n",
      "Model saved as ./models\\resnet18+coordinate_cifar100.pth\n"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def plot_training_results(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs):\n",
    "    # Convert PyTorch tensors to NumPy arrays\n",
    "    train_losses = np.array(train_losses)\n",
    "    train_top1_accs = np.array(train_top1_accs)\n",
    "    train_top5_accs = np.array(train_top5_accs)\n",
    "    val_losses = np.array(val_losses)\n",
    "    val_top1_accs = np.array(val_top1_accs)\n",
    "    val_top5_accs = np.array(val_top5_accs)\n",
    "\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    \n",
    "    # 绘制损失曲线\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(train_losses, label='训练损失')\n",
    "    plt.plot(val_losses, label='验证损失')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('训练/验证损失曲线')\n",
    "    plt.legend()\n",
    "    \n",
    "    # 绘制准确率曲线\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(train_top1_accs, label='训练 Top-1 准确率')\n",
    "    plt.plot(val_top1_accs, label='验证 Top-1 准确率')\n",
    "    plt.plot(train_top5_accs, label='训练 Top-5 准确率')\n",
    "    plt.plot(val_top5_accs, label='验证 Top-5 准确率')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title('训练/验证准确率曲线')\n",
    "    plt.legend()\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# 训练模型\n",
    "def train_model(model, criterion, optimizer, scheduler, num_epochs=25):\n",
    "    since = time.time()\n",
    "\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    " \n",
    "    for epoch in range(num_epochs):\n",
    "        print('第 {}/{} 轮训练'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # 每个epoch都有训练和验证阶段\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # 设置模型为训练模式\n",
    "            else:\n",
    "                model.eval()   # 设置模型为评估模式\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "            running_top1_corrects = 0\n",
    "            running_top5_corrects = 0\n",
    "\n",
    "            # 遍历数据\n",
    "            for inputs, labels in (trainloader if phase == 'train' else testloader):\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # 零化参数梯度\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # 前向传播\n",
    "                # 只在训练阶段追踪历史\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs = model(inputs)\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    loss = criterion(outputs, labels)\n",
    "\n",
    "                    # 计算 top-1 和 top-5 准确率\n",
    "                    _, top5_preds = torch.topk(outputs, 5, dim=1)\n",
    "                    top1_correct = torch.sum(preds == labels.data)\n",
    "                    top5_correct = torch.sum(top5_preds == labels.view(-1, 1))\n",
    "\n",
    "                    # 只有在训练阶段反向传播+优化\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # 统计\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += top1_correct\n",
    "                running_top1_corrects += top1_correct\n",
    "                running_top5_corrects += top5_correct\n",
    "\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "\n",
    "            epoch_loss = running_loss / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top1_acc = running_top1_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top5_acc = running_top5_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "\n",
    "            print('{} 损失: {:.4f} Top-1 准确率: {:.4f} Top-5 准确率: {:.4f}'.format(\n",
    "                phase, epoch_loss, epoch_top1_acc, epoch_top5_acc))\n",
    "            \n",
    "            # 保存损失和准确率\n",
    "            if phase == 'train':\n",
    "                train_losses.append(epoch_loss)\n",
    "                train_top1_accs.append(epoch_top1_acc)\n",
    "                train_top5_accs.append(epoch_top5_acc)\n",
    "            else:\n",
    "                val_losses.append(epoch_loss)\n",
    "                val_top1_accs.append(epoch_top1_acc)\n",
    "                val_top5_accs.append(epoch_top5_acc)\n",
    "\n",
    "            # 深度复制模型\n",
    "            if phase == 'val' and epoch_top1_acc > best_acc:\n",
    "                best_acc = epoch_top1_acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "        print()\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print('训练完成，耗时 {:.0f}m {:.0f}s'.format(\n",
    "        time_elapsed // 60, time_elapsed % 60))\n",
    "    print('最佳验证集准确率: {:4f}'.format(best_acc))\n",
    "\n",
    "    # 加载最佳模型权重\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    \n",
    "    # 绘制训练结果曲线\n",
    "\n",
    "    return model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs)\n",
    "\n",
    "   \n",
    "train_losses = []\n",
    "train_top1_accs = []\n",
    "train_top5_accs = []\n",
    "val_losses = []\n",
    "val_top1_accs = []\n",
    "val_top5_accs = []\n",
    "\n",
    "# 训练模型\n",
    "model_ft = train_model(net, criterion, optimizer, scheduler, num_epochs=25)\n",
    "\n",
    "# 保存模型\n",
    "import os\n",
    "\n",
    "# 确定要保存模型的目录\n",
    "model_dir = './models'\n",
    "os.makedirs(model_dir, exist_ok=True)  # 如果目录不存在，则创建\n",
    "\n",
    "# 保存模型到指定目录\n",
    "model_path = os.path.join(model_dir, 'resnet18+coordinate_cifar100.pth')\n",
    "torch.save(model_ft.state_dict(), model_path)\n",
    "print(f'Model saved as {model_path}')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 绘制训练结果曲线\n",
    "plot_training_results(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
